#pragma once
//a generalized tree
//every node connect to father, siblings and sons

namespace zzz{
template<typename T>
struct TreeNode
{
public:
  T v;
  TreeNode<T> *father;
  vector<TreeNode<T> *> sons;
  int height;

  TreeNode();
  ~TreeNode();

  void AddSon(TreeNode<T> *son);
  bool DeleteSon(TreeNode<T> *son);
  bool DeleteSon(int ison);

  TreeNode<T> *GetFather();
  TreeNode<T> *GetNextSon(TreeNode<T> *son=NULL);
  vector<TreeNode<T> *> GetSons();
  vector<TreeNode<T> *> GetSiblings();
  vector<TreeNode<T> *> GetSameHeight();
  vector<TreeNode<T> *> GetSameHeightSons(int h);

  void GetSiblingsHelperUp(vector<TreeNode<T> *> &result, int h);
  void GetSiblingsHelperDown(vector<TreeNode<T> *> &result, int h);
};

template<typename T>
vector<TreeNode<T> *> zzz::TreeNode<T>::GetSameHeightSons(int h)
{
  vector<TreeNode<T> *> ret;
  GetSiblingsHelperDown(ret, h);
  return ret;
}
template<typename T>
void zzz::TreeNode<T>::GetSiblingsHelperDown(vector<TreeNode<T> *> &result, int h)
{
  if (height==h-1) result.insert(result.end(),sons.begin(),sons.end());
  else
  {
    for (size_t i=0; i<sons.size(); i++)
      sons[i]->GetSiblingsHelperDown(result, h);
  }
}
template<typename T>
void zzz::TreeNode<T>::GetSiblingsHelperUp(vector<TreeNode<T> *> &result, int h)
{
  if (father) father->GetSiblingsHelperUp(result, h);
  vector<TreeNode<T> *> mysiblings=GetSiblings();
  for (size_t i=0; i<mysiblings.size(); i++)
    mysiblings[i]->GetSiblingsHelperDown(result, h);
}
template<typename T>
vector<TreeNode<T> *> zzz::TreeNode<T>::GetSameHeight()
{
  vector<TreeNode<T> *> ret=GetSiblings();
  if (father!=NULL)
    father->GetSiblingsHelperUp(ret,height);
  return ret;
}
template<typename T>
vector<TreeNode<T> *> zzz::TreeNode<T>::GetSiblings()
{
  vector<TreeNode<T> *> ret;
  if (father==NULL) return ret;
  else
  {
    ret=father->GetSons();
    for (size_t i=0; i<ret.size(); i++)
      if (ret[i]==this) 
      {
        ret.erase(ret.begin()+i);
        break;
      }
    return ret;
  }
}
template<typename T>
vector<TreeNode<T> *> zzz::TreeNode<T>::GetSons()
{
  return sons;
}
template<typename T>
TreeNode<T> * zzz::TreeNode<T>::GetNextSon(TreeNode<T> *son/*=NULL*/)
{
  if (son==NULL)
  {
    if (!sons.empty()) return sons[0];
    else return NULL;
  }
  else
  {
    for (size_t i=0; i<sons.size(); i++)
      if (sons[i]==son)
      {
        if (i!=sons.size()-1) return sons[i+1];
        else return NULL;
      }
  }
}
template<typename T>
TreeNode<T> * zzz::TreeNode<T>::GetFather()
{
  return father;
}
template<typename T>
zzz::TreeNode<T>::~TreeNode()
{
  for (size_t i=0; i<sons.size(); i++)
    delete sons[i];
  sons.clear();
}
template<typename T>
zzz::TreeNode<T>::TreeNode()
:father(NULL),height(-1)
{}

template<typename T>
bool zzz::TreeNode<T>::DeleteSon(int ison)
{
  if (i<=(int)sons.size()-1)
  {
    delete sons[i];
    return true;
  }
  return false;
}

template<typename T>
bool zzz::TreeNode<T>::DeleteSon(TreeNode<T> *son)
{
  for (size_t i=0; i<sons.size(); i++)
    if (sons[i]==son)
    {
      delete sons[i];
      return true;
    }
  return false;
}

template<typename T>
void zzz::TreeNode<T>::AddSon(TreeNode<T> *son)
{
  sons.push_back(son);
  son->father=this;
  son->height=height+1;
}

template<typename T>
class Tree
{
public:
  Tree(){head_.height=0;}
  TreeNode<T> *GetHead(){return &head_;}
  TreeNode<T> head_;
};

}